#include "NonLinearFitting.h"

using namespace ypl;

/*******************************************************
 * 类实现：非线性拟合 class NonLinearFittingting::Model
 *******************************************************/

NonLinearFitting::NonLinearFitting(Model& model, Samples& samples)
    :m_pModel(&model), m_pSamples(&samples)
{}


void NonLinearFitting::fitting(Report &report, const LMPara &lmPara)
{
    // 用 L-M 方法拟合
    LMMethod(report, lmPara);
}



void NonLinearFitting::LMMethod(Report &report, const LMPara &lmPara)
{
    // 算法参数
    Real lambda = lmPara.lambda;
    Real step = lmPara.step;
    Real chiEPS = lmPara.chiEPS;
    int maxIter = lmPara.maxIter;

    const Model& model = *m_pModel;

    // 初始参数
    Vector<Real> &p = report.fittedPara;
    model.genParaInBounds(p);

    // 计算初始参数对应的卡方值
    Real chiSq = chiSqure(p);

    int iterCount = 0;
    bool stop = false;
    while (!stop)
    {
        Matrix<Real> alpha;
        Vector<Real> beta;
        calcAlphaBeta(p, alpha, beta);

        // 矩阵 alpha 对角元加上 lambda
        for (int i = 0; i < alpha.rowSize(); ++i)
            alpha[i][i] += lambda;

        // 求解线性方程组 alpha*delta = beta
        Vector<Real> delta;
        LinearEquation::gaussElim(alpha, beta, delta); // 待调整为“黑箱求解”
        // 新的模型参数
        Vector<Real> aNew = p+delta;
        // 新的参数对应卡方值
        Real chiSqNew = chiSqure(aNew);

        if  (chiSqNew >= chiSq)        // 新参数无改进
        {
            lambda *= step;
        }
        else if (chiSqNew < chiSq)   // 新参数有改进
        {
            if (chiSq - chiSqNew < chiEPS)   // 改进较小，认为收敛
            {
                stop = true;
            }
            if (! model.outOfBounds(aNew))
            {
                chiSq = chiSqNew;
                p = aNew;  // 用新参数替换前一步得到的参数
            }

            lambda /= step;
        }

        if (++iterCount >= maxIter)
            stop = true;
    } // while

    // 生成报告
    report.iterCount = iterCount;
    report.chiSquare = chiSq;
    report.rmsError = sqrt(chiSq / m_pSamples->size());
}

/*! @brief 计算卡方函数的曲率矩阵（Hessian 矩阵）与梯度 */
void  NonLinearFittingting::calcAlphaBeta(const Vector<Real>& para,
                                      Matrix<Real>& alpha, Vector<Real>& beta)
{
    Model& y = *m_pModel; // 带参模型
    Integral paraDim = y.paraDim();   // 待参模型参数维数

    if (para.dim() != paraDim)     // 输入参数维数与模型参数维数不一致
        return;

    alpha = Matrix<Real>(paraDim, paraDim);
    beta = Vector<Real>(paraDim);

    const Vector<Real>& p = para;    // 参数值
    Samples& s = *m_pSamples;          // 样本

    Integral sampleSize = s.size();   // 样本数
    Vector<Real> yV(sampleSize); // 与采样点对应的模型计算值
    std::vector<Vector<Real> > vGradient(sampleSize); // 与各采样点对应的梯度
    for (Integral i = 0; i < sampleSize; ++i)
    {
        yV[i] = y(s[i].x(), p);     // y(xi; p)
        vGradient[i] = y.dfdp(s[i].x(), p);   // dy(xi,p)/dp
    }

    bool useDefaultSD = s.useDefaultSD();   // 是否采用默认的标准差 1
    if (useDefaultSD)
    {
        for (Integral k = 0; k < paraDim; ++k)
        {
            // 计算 beta[paraDim]
            Real betaK = 0;   // beta_k
            for (Integral i = 0; i < sampleSize; ++i)
            {
                betaK += (s[i].y() - yV[i]) * vGradient[i][k];
            }
            beta[k] = betaK;
            // 计算 alpha[paraDim][paraDim]
            for (int l = 0; l < paraDim; ++l)
            {
                Real alphaKL = 0; // alpha_kl
                for (Integral i = 0; i < sampleSize; ++i)
                {
                    alphaKL += vGradient[i][k] * vGradient[i][l];
                }
                alpha[k][l] = alphaKL;
            }

        }
    }
    else
    {
        for (Integral k = 0; k < paraDim; ++k)
        {
            // 计算 beta[paraDim]
            Real betaK = 0;   // beta_k
            for (Integral i = 0; i < sampleSize; ++i)
            {
                Real sdi = s[i].sd();
                betaK += (s[i].y() - yV[i]) * vGradient[i][k] / (sdi*sdi);
            }
            beta[k] = betaK;
            // 计算 alpha[paraDim][paraDim]
            for (int l = 0; l < paraDim; ++l)
            {
                Real alphaKL = 0; // alpha_kl
                for (Integral i = 0; i < sampleSize; ++i)
                {
                    Real sdi = s[i].sd();
                    alphaKL += vGradient[i][k] * vGradient[i][l] / (sdi*sdi);
                }
                alpha[k][l] = alphaKL;
            }
        }
    }

}

/*! @brief 给定参数下模型的卡方值 */
Real NonLinearFittingting::chiSqure(const Vector<Real>& para)
{
    const Vector<Real> &p = para;
    Model &y = *m_pModel;  // 带参数模型 y(x; p)
    Samples &s = *m_pSamples;

    Real chiSq = 0.0;
    Integral sampleSize = s.size();
    bool useDefaultSD = s.useDefaultSD(); // 是否使用默认标准差
    if (useDefaultSD)
    {
        for (Integral i = 0; i < sampleSize; ++i)   // 默认标准差为 1.0
        {
            Real tmp = s[i].y() - y(s[i].x(), p);
            chiSq += tmp*tmp;
        }
    }
    else
    {
        for (Integral i = 0; i < sampleSize; ++i)
        {
            Real sdi = s[i].sd(); // 标准差
            Real tmp = s[i].y() - y(s[i].x(), p) / sdi;
            chiSq += tmp*tmp;
        }
    }
    return chiSq;
}

